import os
import pickle
import numpy as np
from scipy.ndimage import distance_transform_edt as distance
from skimage import segmentation as skimage_seg
import torch
from torch.utils.data.sampler import Sampler
import  cv2 as cv




def img_to_sdf(img_gt, out_shape):
    """
    compute the signed distance map of binary mask
    input: segmentation, shape = (batch_size, x, y, z)
    output: the Signed Distance Map (SDM)
    sdf(x) = 0; x in segmentation boundary
             -inf|x-y|; x in segmentation
             +inf|x-y|; x out of segmentation
    normalize sdf to [-1,1]
    """
    #  out_shape = (batchsize,h,w)  ,,img_gt.shape=[4, 1, 256, 256]

    img_gt = img_gt.astype(np.uint8)
    normalized_sdf = np.zeros(out_shape)
    gt_sdf = np.zeros(out_shape)
    for b in range(out_shape[0]): # batch size
        # print("b==",b)
        # print("img_gt.unique==",np.unique(img_gt))
        # print("img.shape",img_gt.shape)
        # cv.imshow("B", img_gt[0,:,:,0])
        # cv.waitKey(-1)
        # exit()
        # posmask.shape = [1,256,256]
        posmask = img_gt[b].astype(np.bool)
        if posmask.any():
            # [True,False] --> [False,True]
            negmask = ~posmask
            # distance 即 distance_transform_edt是scipy库里的一个函数，用于距离转换，计算图像中非零点到最近背景点（即0）的距离。
            # print("posmask==",posmask)
            posdis = distance(posmask)
            # print("posdis==",posdis)

            # print("negmask==", negmask)
            negdis = distance(negmask)
            # print("negdis==", negdis)

            # bool 转 uint8,true和1表示边界
            boundary = skimage_seg.find_boundaries(posmask, mode='inner').astype(np.uint8)

            # print("boundary.unique=", np.unique(boundary))
            # print("posdis.unique=", np.unique(posdis))
            #
            # print("negdis.unique=", np.unique(negdis))
            # print("boundary.unique=", np.unique(boundary))

            # 归一化到 [-1,1]
            # print("-->",posdis.shape)  #(1, 256, 256)
            # with open(r"G:\Desktop\program\myProgram\ccnet\testing\1.txt","w") as f:
            #     for i in range(len(posdis[0])):
            #         f.write(str(posdis[0][i])+"\n")


            # print("--",np.squeeze(posdis,axis=0).shape)
            tmp = np.squeeze(posdis,axis=0)
            # print("-->",np.where(tmp==2))
            xx,yy = np.where(tmp==2)
            # print("-->",tmp[240:256,240:256])
            # print("val",posdis[0][xx[0]][yy[0]])

            # negdis:背景距离边界/差值    posdis:分割内距离边界/差值
            # sdf对应背景为正，，对应分割为负
            sdf = (negdis-np.min(negdis))/(np.max(negdis)-np.min(negdis)) - (posdis-np.min(posdis))/(np.max(posdis)-np.min(posdis))

            solid_sdf = negdis - posdis

            lis = []
            for i,j in zip(xx,yy):
                # print("val", posdis[0][i][j],end="-----------")
                # print("sdf", sdf[0][xx[i]][yy[j]])
                lis.append(sdf[0][i][j])

            # print("max:{0}  min:{1}".format(max(lis),min(lis)))
            # 1--> sdf -0.037011660509880265
            # val 1.4142135623730951-----------sdf -0.052342392259021375
            #
            # 可能posmask里值为1的位置认为是边界
            # print("sdf.unique=",np.unique(sdf))
            # exit()
            # boundary == 1的位置就是边界
            sdf[boundary==1] = 0

            solid_sdf[boundary==1] = 0

            # print("sdf", sdf[0][xx[0]][yy[0]])
            normalized_sdf[b] = sdf
            gt_sdf[b] = solid_sdf
            # exit()
            # assert np.min(sdf) == -1.0, print(np.min(posdis), np.max(posdis), np.min(negdis), np.max(negdis))
            # assert np.max(sdf) ==  1.0, print(np.min(posdis), np.min(negdis), np.max(posdis), np.max(negdis))

    return normalized_sdf,gt_sdf


def sigmoid_rampup(current, rampup_length):
    """Exponential rampup from https://arxiv.org/abs/1610.02242"""
    if rampup_length == 0:
        return 1.0
    else:
        # clip这个函数将将数组current中的元素限制在a_min, a_max之间，大于a_max的就使得它等于 a_max，小于a_min,的就使得它等于a_min
        current = np.clip(current, 0.0, rampup_length)
        phase = 1.0 - current / rampup_length
        return float(np.exp(-5.0 * phase * phase))


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        # self.avg = 0
        # self.sum = 0
        # self.count = 0

    def update(self, val, n=1):
        self.val = val + self.val
        # self.sum += val * n
        # self.count += n
        # self.avg = self.sum / self.count






















